﻿#ifndef THREADPOOL_H
#define THREADPOOL_H

#include <queue>
#include <memory>
#include <atomic>
#include <mutex>
#include <functional>
#include <condition_variable>
#include <thread>
#include <unordered_map>
#include <future>
#include <iostream>

const int TASK_MAX_THRESHOLD = 2;
const int THREAD_MAX_THRESHOLD = 10;
const int THREAD_MAX_IDLE_TIME = 1;    // 秒

// 线程池支持的模式
enum class PoolMode
{
    MODE_FIXED,     // 固定数量
    MODE_CACHED,    // 数量可动态增长 适合小而短的任务
};

class ThreadPool
{
public:
    ThreadPool()
        : initThreadSize_(4)
        , idleThreadCount_(0)
        , curThreadCount_(0)
        , taskSize_(0)
        , taskQueMaxThreshold_(TASK_MAX_THRESHOLD)
        , threadCountMaxThreshold_(THREAD_MAX_THRESHOLD)
        , poolMode_(PoolMode::MODE_FIXED)
        , isRunning_(false)  // 初始为运行状态
    {
    }
    ~ThreadPool()
    {
        // 设置停止标志
        isRunning_.store(false);

        // 通知所有线程解除阻塞 以便退出
        notEmpty_.notify_all();

        // 回收线程资源
        {
            std::unique_lock<std::mutex> mapLock(threadMapMutex_);
            for (auto& [id, thread] : threads_)
            {
                if (thread.joinable())
                {
                    thread.join();
                }
            }
        }

        threads_.clear(); // 清空容器
    }

    // 设置线程池模式
    void setMode(PoolMode mode)
    {
        if (checkRunningState())
        {
            return;
        }

        poolMode_ = mode;
    }

    // 开启线程池
    void start(int initThreadCount = std::thread::hardware_concurrency())
    {
        // 设置线程池的运行状态
        isRunning_ = true;

        // 记录初始线程数量
        initThreadSize_ = initThreadCount;

        // 创建线程对象
        std::unique_lock<std::mutex> threadsLock(threadMapMutex_);
        for (int i = 0; i < initThreadSize_; ++i)
        {
            // 创建线程对象的时候 把线程函数传入
            std::thread t([this] { threadFunc(); });
            threads_.emplace(t.get_id(), std::move(t));
            ++curThreadCount_;
            ++idleThreadCount_;
        }
    }

    // 设置任务队列上限阈值
    void setTaskQueMaxThreashold(int threashold)
    {
        if (checkRunningState())
        {
            return;
        }

        taskQueMaxThreshold_ = threashold;
    }

    // 设置线程数量上限阈值
    void setThreadCountMaxThreashold(int threashold)
    {
        if (checkRunningState() || poolMode_ != PoolMode::MODE_CACHED)
        {
            return;
        }

        threadCountMaxThreshold_ = threashold;
    }

    // 给线程池提交任务
    template<typename Func, typename... Args>
    auto submitTask(Func&& func, Args&&... args) -> std::future<decltype(func(args...))>
    {
        // 打包任务 放入任务队列
        using ReturnType = decltype(func(args...));
        // 例子：packaged_task<int(int, int)> task(std::bind(sum, 11, 22));
        // 这里不能直接创建task对象。task是submitTask函数内的局部变量，当函数返回时，
        // task被销毁，但任务队列中的lambda仍持有其引用。
        // 线程执行任务时，实际调用的是一个已销毁的packaged_task对象，导致未定义行为。
        // 所以需要用共享指针延长其生命周期
        auto task = std::make_shared<std::packaged_task<ReturnType()>>(
            std::bind(std::forward<Func>(func), std::forward<Args>(args)... )
        );
        std::future<ReturnType> result = task->get_future();

        // 加锁
        std::unique_lock<std::mutex> lock(taskQueMtx_);

        // 任务队列满了 需要等待
        // 用户提交任务 最长不能阻塞超过1s 否则判断提交任务失败 返回
        if (!notFull_.wait_for(lock, std::chrono::seconds(1),
            [this] {return taskQue_.size() < (size_t)taskQueMaxThreshold_; }))
        {
            // 任务队列满了 1s内没有空余位置
            std::cerr << "task queue is full, submit task failed!" << std::endl;
            auto rtTask = std::make_shared<std::packaged_task<ReturnType()>>(
                []() ->ReturnType { return ReturnType(); }
            );
            (*rtTask)();
            return rtTask->get_future();
        }

        // 任务队列不满 直接添加任务
        // 这里需要按值引用 复制一份共享指针
        taskQue_.emplace([task]() { (*task)(); });
        // 此时的引用计数是2
        //std::cout << "引用计数：" << task.use_count() << std::endl;
        ++taskSize_;

        // 用Notempty通知一个线程
        notEmpty_.notify_all();

        // cached模式 需要根据任务数量和空闲线程的数量，判断是否需要创建新的线程出来
        if (poolMode_ == PoolMode::MODE_CACHED
            && taskSize_ > idleThreadCount_
            && curThreadCount_ <= threadCountMaxThreshold_)
        {
            std::thread t([this] { threadFunc(); });
            std::cout << "新增了线程：" << t.get_id() << std::endl;

            std::unique_lock<std::mutex> mapLock(threadMapMutex_);
            threads_.emplace(t.get_id(), std::move(t));
            mapLock.unlock();

            ++curThreadCount_;
            ++idleThreadCount_;
        }

        // 返回任务的Result对象
        return result;
    }

    ThreadPool(const ThreadPool&) = delete;
    ThreadPool& operator=(const ThreadPool&) = delete;

private:
    // 定义线程函数
    void threadFunc()
    {
        auto lastTime = std::chrono::high_resolution_clock().now();

        while (true)
        {
            // 先获取锁
            std::unique_lock<std::mutex> lock(taskQueMtx_);

            std::cout << "tid:" << std::this_thread::get_id() <<
                "尝试获取任务" << std::endl;

            // cached模式下，有可能已经创建了很多的线程，但是空闲时间超过60s，
            // 应该把多余的线程/结束回收掉（超过initThreadsize_数量的线程要进行回收)
            // 当前时间 - 上一次线程执行的时间 >= THREAD_MAX_IDLE_TIME
            // 锁 + isRunning_ 双重判断 避免死锁
            while (taskQue_.empty())
            {
                if (!isRunning_.load())
                {
                    auto threadID = std::this_thread::get_id();
                    std::cout << "已结束线程：" << threadID << std::endl;
                    return;
                }

                if (poolMode_ == PoolMode::MODE_CACHED)
                {
                    // 条件变量超时才往下进行
                    if (std::cv_status::timeout !=
                        notEmpty_.wait_for(lock, std::chrono::seconds(1)))
                    {
                        continue;
                    }

                    auto nowTime = std::chrono::high_resolution_clock().now();
                    auto delta = std::chrono::duration_cast<std::chrono::seconds>(
                        nowTime - lastTime);
                    // 空闲时间短或当前线程没超过初始线程 就不回收
                    if (delta.count() < THREAD_MAX_IDLE_TIME
                        || curThreadCount_ <= initThreadSize_)
                    {
                        continue;
                    }

                    // 开始回收当前线程
                    // 把线程对象从线程列表容器中删除
                    auto threadID = std::this_thread::get_id();
                    std::cout << "超时线程的ID：" << threadID << std::endl;

                    // 在需要操作容器的地方加锁
                    std::unique_lock<std::mutex> maplock(threadMapMutex_);
                    auto it = threads_.find(threadID);
                    if (it != threads_.end())
                    {
                        if (it->second.joinable())
                        {
                            // 如果在threadFunc里调用自身的join()
                            // 会导致线程无限等待自己结束，造成死锁
                            // 所以使用detach()替代join()，使线程分离，允许其自然结束而不阻塞自身。
                            it->second.detach();
                        }
                        threads_.erase(it);
                        std::cout << "已移除空闲线程：" << threadID << std::endl;
                    }
                    maplock.unlock();

                    // 记录线程数量的相关变量的值修改 
                    --curThreadCount_;
                    --idleThreadCount_;

                    return; // 退出运行
                }
                else
                {
                    // 等待notEmpty条件
                    notEmpty_.wait(lock);
                }
            }

            --idleThreadCount_;

            // 从任务队列中取一个任务出来
            auto task = taskQue_.front();
            taskQue_.pop();
            --taskSize_;

            std::cout << "tid:" << std::this_thread::get_id() <<
                "获取任务成功" << std::endl;

            // 如果还有剩余任务 继续通知其他的线程执行任务
            if (taskQue_.size() > 0)
            {
                notEmpty_.notify_all();
            }

            // 释放锁
            lock.unlock();
            notFull_.notify_all();

            // 当前线程负责执行这个任务
            if (task)
            {
                task();
            }

            ++idleThreadCount_;

            // 更新时间
            lastTime = std::chrono::high_resolution_clock().now();
        }
    }

    // 检查pool的运行状态
    bool checkRunningState() const
    {
        return isRunning_;
    }

private:
    std::unordered_map<std::thread::id, std::thread> threads_;  // 线程列表
    std::mutex threadMapMutex_;         // 保护 threads_ 的互斥锁

    int initThreadSize_;                // 初始线程数量
    int threadCountMaxThreshold_;       // 线程数量上限
    std::atomic_int curThreadCount_;    // 线程数量
    std::atomic_int idleThreadCount_;   // 空闲线程的数量

    // 由于无法事先知道被提交任务的返回值类型
    // 因此可用void返回值类型的lambda表达式充当中间层
    // 使用时传入对应任务的副本并执行即可：
    //      taskQue_.emplace([task]() { (*task)(); });
    using Task = std::function<void()>;
    std::queue<Task> taskQue_;          // 任务队列
    std::mutex taskQueMtx_;             // 保护 taskQue_ 的互斥锁
    std::atomic_int taskSize_;          // 任务数量
    int taskQueMaxThreshold_;           // 任务数量上限
    std::condition_variable notFull_;   // 任务队列不满
    std::condition_variable notEmpty_;  // 任务队列不空

    PoolMode poolMode_;             // 当前线程池的工作模式
    std::atomic_bool isRunning_;    // 控制线程是否继续运行
};

#endif // !THREADPOOL_H
